# kline_utils.py
from Ashare import get_price
import pandas as pd
import os



def save_kline_data(stock_code, frequency = '15m', count=100):
    output = fr'D:\work\code\clark\gitee\big_a\datas\kline_now\{stock_code}\{frequency}.json'
    # 获取新数据
    df = get_price(stock_code, frequency=frequency, count=count)
    df = df.reset_index()
    if 'index' in df.columns:
        df.rename(columns={'index': 'date'}, inplace=True)
    else:
        time_col = df.columns[0]
        df.rename(columns={time_col: 'date'}, inplace=True)
    df['date'] = pd.to_datetime(df['date'])
    df['date'] = df['date'].dt.strftime('%Y-%m-%d %H:%M:%S')

    # 创建目录（如果不存在）
    output_dir = os.path.dirname(output)
    if not os.path.exists(output_dir):
        try:
            os.makedirs(output_dir, exist_ok=True)
            print(f"✅ 目录创建成功：{output_dir}")
        except OSError as e:
            print(f"❌ 目录创建失败：{e}")
            return  # 可以选择抛出异常，这里直接返回
    else:
        print(f"ⓘ 目录已存在：{output_dir}")

    # 如果文件存在，则读取旧数据并合并去重；否则直接保存
    if os.path.exists(output):
        old_df = pd.read_json(output, orient='records')
        # 统一转换为Timestamp类型再排序
        old_df['date'] = pd.to_datetime(old_df['date'])
        df['date'] = pd.to_datetime(df['date'])
        # 合并并去重（新数据覆盖旧数据）
        combined = pd.concat([old_df, df]).drop_duplicates('date', keep='last')
        # 排序
        combined = combined.sort_values('date').reset_index(drop=True)
        # 转回字符串格式
        combined['date'] = combined['date'].dt.strftime('%Y-%m-%d %H:%M:%S')
        # 保存
        combined.to_json(output, orient='records', indent=4)
        added_count = len(combined) - len(old_df)  # 计算新增记录数（可能是正数或0）
        print(f"✅ 数据已追加更新至 {output} (新增 {added_count} 条 共 {len(combined)} 条)")
    else:
        df.to_json(output, orient='records', indent=4)
        print(f"✅ 新文件已创建于 {output} (保存 {len(df)} 条)")